import azure.functions as func
import azure.durable_functions as df

from azure.monitor.opentelemetry import configure_azure_monitor
from opentelemetry import trace
from opentelemetry.trace import SpanKind
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

configure_azure_monitor()
tracer = trace.get_tracer(__name__)

myApp = df.DFApp(http_auth_level=func.AuthLevel.ANONYMOUS)

@myApp.route(route="entity_add_value")
@myApp.durable_client_input(client_name="client")
async def http_add(req: func.HttpRequest, client, context: func.Context):
    # Extract the traceparent and tracestate from the context
    parent = {
        "traceparent": context.trace_context.trace_parent,
        "tracestate": context.trace_context.trace_state
    }
    parent_context = TraceContextTextMapPropagator().extract(parent)

    # Create a new span with the extracted parent context
    with tracer.start_as_current_span(
            "myCustomSpan",
            context=parent_context,
            kind=SpanKind.CLIENT) as span:
        
        span.set_attribute("az.namespace", "app")

        entityId = df.EntityId("Counter", "myCounter")
        await client.signal_entity(entityId, "add", 1)
        response = func.HttpResponse("Done", status_code=200)
    return response

@myApp.route(route="orchestrator")
@myApp.durable_client_input(client_name="client")
async def http_start(req: func.HttpRequest, client, context: func.Context):
    # Extract the traceparent and tracestate from the context
    parent = {
        "traceparent": context.trace_context.trace_parent,
        "tracestate": context.trace_context.trace_state
    }
    parent_context = TraceContextTextMapPropagator().extract(parent)

    # Create a new span with the extracted parent context
    with tracer.start_as_current_span(
            "myCustomSpan",
            context=parent_context,
            kind=SpanKind.CLIENT) as span:
        
        span.set_attribute("az.namespace", "app")

        instance_id = await client.start_new('orchestrator')
        response = client.create_check_status_response(req, instance_id)
        return response

@myApp.orchestration_trigger(context_name="context")
def orchestrator(context: df.DurableOrchestrationContext):
    entityId = df.EntityId("Counter", "myCounter")
    context.signal_entity(entityId, "add", 3)
    state = yield context.call_entity(entityId, "get")
    return state